import torch
import torch.nn as nn
from .CVT import *

__all__ = ["ResCVT"]


class ResCVT(nn.Module):
    def __init__(
            self,
            *,
            num_classes,
            s1_emb_dim=64,
            s1_emb_kernel=7,
            s1_emb_stride=4,
            s1_proj_kernel=3,
            s1_kv_proj_stride=2,
            s1_heads=1,
            s1_depth=1,
            s1_mlp_mult=4,
            s2_emb_dim=192,
            s2_emb_kernel=3,
            s2_emb_stride=2,
            s2_proj_kernel=3,
            s2_kv_proj_stride=2,
            s2_heads=3,
            s2_depth=2,
            s2_mlp_mult=4,
            s3_emb_dim=384,
            s3_emb_kernel=3,
            s3_emb_stride=2,
            s3_proj_kernel=3,
            s3_kv_proj_stride=2,
            s3_heads=6,
            s3_depth=10,
            s3_mlp_mult=4,
            dropout=0.
    ):
        super().__init__()
        kwargs = dict(locals())

        dim = 3
        layers = []

        for prefix in ('s1', 's2', 's3'):
            config, kwargs = group_by_key_prefix_and_remove_prefix(f'{prefix}_', kwargs)

            layers.append(nn.Sequential(
                nn.Conv2d(dim, config['emb_dim'], kernel_size=config['emb_kernel'], padding=(config['emb_kernel'] // 2),
                          stride=config['emb_stride']),
                LayerNorm(config['emb_dim']),
                Transformer(dim=config['emb_dim'], proj_kernel=config['proj_kernel'],
                            kv_proj_stride=config['kv_proj_stride'], depth=config['depth'], heads=config['heads'],
                            mlp_mult=config['mlp_mult'], dropout=dropout)
            ))

            dim = config['emb_dim']

        self.layers = nn.Sequential(*layers)

        self.to_logits = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            Rearrange('... () () -> ...'),
            nn.Linear(dim, num_classes)
        )

    def forward(self, x):
        out = x
        for layer in self.layers:
            out = x + layer(out)
        return self.to_logits(out)